{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Whole-body manipulation\n",
    "The objective of this exercise is to reach multiple targets while keeping balance in the Talos humanoid robot.\n",
    "<img src=\"https://pal-robotics.com/wp-content/uploads/2024/04/TALOS_datasheets-rayos-768x1349.webp\" alt=\"drawing\" width=\"250\"/>\n",
    "\n",
    "\n",
    "This exercise focuses on a multi-contact optimal control problem of the form:\n",
    "\n",
    "\\begin{equation}\\nonumber\n",
    "\t\\begin{aligned}\n",
    "\t\t\\min_{\\mathbf{x}_s,\\mathbf{u}_s}\n",
    "\t\t&\\hspace{-2.em}\n",
    "\t\t& & \\hspace{-0.75em}l_N(\\mathbf{x}_{N})+\\sum_{k=0}^{N-1} \\int_{t_k}^{t_k+\\Delta t_k}\\hspace{-2.em} l_k(\\mathbf{x}_k,\\mathbf{u}_k)dt \\hspace{-8.em}&\\\\\n",
    "\t\t& \\hspace{-1em}\\textrm{s.t.}\n",
    "\t\t& & \\mathbf{q}_{k+1} = \\mathbf{q}_k \\oplus \\int_{t_k}^{t_k+\\Delta t_k}\\hspace{-2.em}\\mathbf{v}_{k+1}\\,dt, &\\textrm{(integrator)}\\\\\n",
    "\t\t& & & \\mathbf{v}_{k+1} = \\mathbf{v}_k + \\int_{t_k}^{t_k+\\Delta t_k}\\hspace{-2.em}\\mathbf{\\dot{v}}_k\\,dt, &\\\\\n",
    "\t\t& & & \\hspace{-1em}\\left[\\begin{matrix}\\mathbf{\\dot{v}}_k \\\\ -\\boldsymbol{\\lambda}_k\\end{matrix}\\right] =\n",
    "\t\t\\left[\\begin{matrix}\\mathbf{M} & \\mathbf{J}^{\\top}_c \\\\ {\\mathbf{J}_{c}} & \\mathbf{0} \\end{matrix}\\right]^{-1}\n",
    "\t\t\\left[\\begin{matrix}\\boldsymbol{\\tau}_b \\\\ -\\mathbf{a}_0 \\\\\\end{matrix}\\right], &\\textrm{(contact dynamics)}\\\\\n",
    "\t\t& & & \\mathbf{R}\\boldsymbol{\\lambda}_{\\mathcal{C}(k)} \\leq \\mathbf{\\mathbf{r}}, &\\textrm{(friction-cone)}\\\\\n",
    "\t\t\t& & & \\mathbf{\\underline{x}} \\leq \\mathbf{x}_k \\leq \\mathbf{\\bar{x}}, &\\textrm{(state bounds)}\n",
    "\t\t\\end{aligned}\n",
    "\\end{equation}\n",
    "\n",
    "where $l_i(\\mathbf{x}_i, \\mathbf{u}_i) = w_{hand}\\|\\log{(\\mathbf{p}_{\\mathcal{G}(k)}(\\mathbf{q}_k)^{-1} \\mathbf{^oM}_{\\mathbf{f}_{\\mathcal{G}(k)}})}\\| + w_{xreg}\\|\\mathbf{x} - \\mathbf{x}_0\\|_{Q} + w_{ureg}\\|\\mathbf{u}\\|_{R}$. Note that (1) the first term is the hand placement cost and (2) the terminal cost does not include the control regularization term.\n",
    "\n",
    "Below there is a basic example that defines the above problem for reaching one target. Later, you will have to build the problem on top of it.\n",
    "\n",
    "Without no more preamble, let's first declare the robot model and the foot and hand names!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import crocoddyl\n",
    "import example_robot_data\n",
    "import numpy as np\n",
    "import pinocchio as pin\n",
    "\n",
    "# Load robot\n",
    "robot = example_robot_data.load(\"talos\")\n",
    "rmodel = robot.model\n",
    "q0 = rmodel.referenceConfigurations[\"half_sitting\"]\n",
    "x0 = np.concatenate([q0, np.zeros(rmodel.nv)])\n",
    "\n",
    "# Declaring the foot and hand names\n",
    "rf_name = \"right_sole_link\"\n",
    "lf_name = \"left_sole_link\"\n",
    "lh_name = \"gripper_left_joint\"\n",
    "\n",
    "# Getting the frame ids\n",
    "rf_id = rmodel.getFrameId(rf_name)\n",
    "lf_id = rmodel.getFrameId(lf_name)\n",
    "lh_id = rmodel.getFrameId(lh_name)\n",
    "\n",
    "# Define the robot's state and actuation\n",
    "state = crocoddyl.StateMultibody(rmodel)\n",
    "actuation = crocoddyl.ActuationModelFloatingBase(state)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With the following function, we could build a differential action model giving a desired hand target.\n",
    "The function builds a double-support contact phase and defines a hand-placement task. The cost function also includes:\n",
    " - state and control regularization terms\n",
    " - state limits penalization\n",
    " - friction cone penalization\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def createActionModel(target = None):\n",
    "    # Creating a double-support contact (feet support)\n",
    "    contacts = crocoddyl.ContactModelMultiple(state, actuation.nu)\n",
    "    lf_contact = crocoddyl.ContactModel6D(\n",
    "        state,\n",
    "        lf_id,\n",
    "        pin.SE3.Identity(),\n",
    "        pin.LOCAL_WORLD_ALIGNED,\n",
    "        actuation.nu,\n",
    "        np.array([0, 40.]),\n",
    "    )\n",
    "    rf_contact = crocoddyl.ContactModel6D(\n",
    "        state,\n",
    "        rf_id,\n",
    "        pin.SE3.Identity(),\n",
    "        pin.LOCAL_WORLD_ALIGNED,\n",
    "        actuation.nu,\n",
    "        np.array([0, 40.]),\n",
    "    )\n",
    "    contacts.addContact(\"lf_contact\", lf_contact)\n",
    "    contacts.addContact(\"rf_contact\", rf_contact)\n",
    "\n",
    "    # Define the cost sum (cost manager)\n",
    "    costs = crocoddyl.CostModelSum(state, actuation.nu)\n",
    "\n",
    "    # Adding the hand-placement cost\n",
    "    if target is not None:\n",
    "        w_hand = np.array([1] * 3 + [0.0001] * 3)\n",
    "        lh_Mref = pin.SE3(np.eye(3), target)\n",
    "        activation_hand = crocoddyl.ActivationModelWeightedQuad(w_hand**2)\n",
    "        lh_cost = crocoddyl.CostModelResidual(\n",
    "            state,\n",
    "            activation_hand,\n",
    "            crocoddyl.ResidualModelFramePlacement(state, lh_id, lh_Mref, actuation.nu),\n",
    "        )\n",
    "        costs.addCost(\"lh_goal\", lh_cost, 1e2)\n",
    "\n",
    "    # Adding state and control regularization terms\n",
    "    w_x = np.array([0] * 3 + [10.0] * 3 + [0.01] * (state.nv - 6) + [10] * state.nv)\n",
    "    activation_xreg = crocoddyl.ActivationModelWeightedQuad(w_x**2)\n",
    "    x_reg_cost = crocoddyl.CostModelResidual(\n",
    "        state, activation_xreg, crocoddyl.ResidualModelState(state, x0, actuation.nu)\n",
    "    )\n",
    "    u_reg_cost = crocoddyl.CostModelResidual(\n",
    "        state, crocoddyl.ResidualModelControl(state, actuation.nu)\n",
    "    )\n",
    "    costs.addCost(\"xReg\", x_reg_cost, 1e-3)\n",
    "    costs.addCost(\"uReg\", u_reg_cost, 1e-4)\n",
    "\n",
    "    # Adding the state limits penalization\n",
    "    x_lb = np.concatenate([state.lb[1 : state.nv + 1], state.lb[-state.nv :]])\n",
    "    x_ub = np.concatenate([state.ub[1 : state.nv + 1], state.ub[-state.nv :]])\n",
    "    activation_xbounds = crocoddyl.ActivationModelQuadraticBarrier(\n",
    "        crocoddyl.ActivationBounds(x_lb, x_ub)\n",
    "    )\n",
    "    x_bounds = crocoddyl.CostModelResidual(\n",
    "        state,\n",
    "        activation_xbounds,\n",
    "        crocoddyl.ResidualModelState(state, actuation.nu),\n",
    "    )\n",
    "    costs.addCost(\"xBounds\", x_bounds, 1.0)\n",
    "\n",
    "    # Adding the friction cone penalization\n",
    "    nsurf, mu = np.identity(3), 0.7\n",
    "    cone = crocoddyl.FrictionCone(nsurf, mu, 4, False)\n",
    "    activation_friction = crocoddyl.ActivationModelQuadraticBarrier(\n",
    "        crocoddyl.ActivationBounds(cone.lb, cone.ub)\n",
    "    )\n",
    "    lf_friction = crocoddyl.CostModelResidual(\n",
    "        state,\n",
    "        activation_friction,\n",
    "        crocoddyl.ResidualModelContactFrictionCone(state, lf_id, cone, actuation.nu),\n",
    "    )\n",
    "    rf_friction = crocoddyl.CostModelResidual(\n",
    "        state,\n",
    "        activation_friction,\n",
    "        crocoddyl.ResidualModelContactFrictionCone(state, rf_id, cone, actuation.nu),\n",
    "    )\n",
    "    costs.addCost(\"lf_friction\", lf_friction, 1e1)\n",
    "    costs.addCost(\"rf_friction\", rf_friction, 1e1)\n",
    "\n",
    "    # Creating the action model\n",
    "    dmodel = crocoddyl.DifferentialActionModelContactFwdDynamics(\n",
    "        state, actuation, contacts, costs\n",
    "    )\n",
    "    return dmodel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, the following function allows us to display the motions and desired targets:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import meshcat.geometry as g\n",
    "\n",
    "def createDisplay(targets):\n",
    "    display = crocoddyl.MeshcatDisplay(robot, 4, 4, False)\n",
    "    for i, target in enumerate(targets):\n",
    "        display.robot.viewer[\"target_\" + str(i)].set_object(g.Sphere(0.05))\n",
    "        display.robot.viewer[\"target_\" + str(i)].set_transform(\n",
    "            np.array(\n",
    "                [\n",
    "                    [1.0, 0.0, 0.0, target[0]],\n",
    "                    [0.0, 1.0, 0.0, target[1]],\n",
    "                    [0.0, 0.0, 1.0, target[2]],\n",
    "                    [0.0, 0.0, 0.0, 1.0],\n",
    "                ]\n",
    "            )\n",
    "        )\n",
    "    return display"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we create an optimal control problem to reach a single target"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DT, N = 5e-2, 20\n",
    "target = np.array([0.4, 0, 1.2])\n",
    "\n",
    "# Creating a running model for the target\n",
    "running_models = [crocoddyl.IntegratedActionModelEuler(createActionModel(target), DT)] * N\n",
    "terminal_model = crocoddyl.IntegratedActionModelEuler(createActionModel(target), 0.0)\n",
    "print(\"Running models:\", running_models[0])\n",
    "print(\"Terminal model:\", terminal_model)\n",
    "\n",
    "\n",
    "# Defining the problem and the solver\n",
    "problem = crocoddyl.ShootingProblem(x0, running_models, terminal_model)\n",
    "fddp = crocoddyl.SolverFDDP(problem)\n",
    "\n",
    "# Creating display\n",
    "display = createDisplay([target])\n",
    "\n",
    "# Adding callbacks to inspect the evolution of the solver (logs are printed in the terminal)\n",
    "fddp.setCallbacks([crocoddyl.CallbackVerbose()])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's solve this problem!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "status = fddp.solve()\n",
    "print(\"\\nProblem solved:\", status, flush=True)\n",
    "print(\"Number of iterations:\", fddp.iter)\n",
    "print(\"Total cost:\", fddp.cost)\n",
    "print(\"Gradient norm:\", fddp.stoppingCriteria())\n",
    "\n",
    "# Embedded in this cell\n",
    "display.robot.viewer.jupyter_cell()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You could display again the final solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display.rate = -1\n",
    "display.freq = 1\n",
    "\n",
    "display.displayFromSolver(fddp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Modifying the example\n",
    "\n",
    "Let's build an optimal control problem to reach 4 targets as described below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "targets = []\n",
    "targets += [np.array([0.4, 0.1, 1.2])]\n",
    "targets += [np.array([0.6, 0.1, 1.2])]\n",
    "targets += [np.array([0.6, -0.1, 1.2])]\n",
    "targets += [np.array([0.4, -0.1, 1.2])]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's display the targets in Meshcat. Do not forget to embed again the display into the jupyter cell"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Embedded in this cell\n",
    "display.robot.viewer.jupyter_cell()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After checking that everything is alright, it's time to build the sequence!\n",
    "Do not forget to create the problem as well :)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And we solve it as before"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the FDDP solver\n",
    "fddp = crocoddyl.SolverFDDP(problem)\n",
    "fddp.setCallbacks([crocoddyl.CallbackVerbose(), crocoddyl.CallbackDisplay(display)])\n",
    "\n",
    "# Solves the problem\n",
    "print(\"Problem solved:\", fddp.solve())\n",
    "print(\"Number of iterations:\", fddp.iter)\n",
    "print(\"Total cost:\", fddp.cost)\n",
    "print(\"Gradient norm:\", fddp.stoppingCriteria())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Do not miss the change to display the motion at the right display speed!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display.rate = -1\n",
    "display.freq = 1\n",
    "display.displayFromSolver(fddp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Same targets with right hand\n",
    "\n",
    "You've learned how to reach 4 targets with the left hand, congratulations!\n",
    "\n",
    "To keep playing within this problem, you should create a new createActionModel to achieve the same task for the right hand."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def createActionModel(target):\n",
    "    # now god is with you xD\n",
    "    # time to show you up!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And here you need to create the problem and solve.\n",
    "Do not forget to display the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "crocoddyl",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
